#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import numpy as np
import cv2
from matplotlib import pyplot as plt
#plt.switch_backend('Agg')
def zifu_segment(path):
    count=0
    img=cv2.imread(path,0)
    _,binary_image=cv2.threshold(img,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 60))
    opened = cv2.morphologyEx(binary_image, cv2.MORPH_OPEN, kernel, iterations=1)
    opened=255-opened
    binary_image=255-binary_image
    contours, _ = cv2.findContours(opened, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # Find contours
    cntsSorted = sorted(contours, key=lambda x: cv2.contourArea(x), reverse=True)  # Sort the contours
    cntsLength = len(cntsSorted)
    images = []
    for idx in range(cntsLength):  # Iterate over the contours
        contour_no = contours[idx]  # Get the first number

        x, y, w, h = cv2.boundingRect(contour_no)  # Get its position and size
        no_tl = (x, y)
        no_br = (x + w, y + h)

        sample_no = binary_image[no_tl[1]:no_br[1], no_tl[0]:no_br[0]]  # Crop the number from the original image

        # Good squared
        offset_reference = max(sample_no.shape[0], sample_no.shape[1])
        pad_offset_height = (offset_reference - sample_no.shape[0]) // 2
        pad_offset_width = (offset_reference - sample_no.shape[1]) // 2
        sample_no = np.pad(sample_no,
                           [(pad_offset_height, pad_offset_height), (pad_offset_width, pad_offset_width)],
                           mode='constant')  # height_b, height_a, width_b, width_a
        # Margin
        pad_height = int(sample_no.shape[0] / 6)
        pad_width = int(sample_no.shape[1] / 6)
        sample_no = np.pad(sample_no, [(pad_height, pad_height), (pad_width, pad_width)],
                           mode='constant')  # height_b, height_a, width_b, width_a

        images.append([x, sample_no])  # Add the image to the list of images and the X position
    images = sorted(images, key=lambda x: x[0])  # Sort the list of images using the X position
    imagesLength = len(images)
    for idx in range(imagesLength):
        plt.subplot(1, imagesLength, idx + 1),
        kernel1=np.ones((4,4),np.uint8)
        plt.imshow(cv2.dilate(images[idx][1],kernel1,iterations=1), 'gray')  # Add every number to the plot
        plt.xticks([]), plt.yticks([])  # Delete the axis
        image = cv2.resize(cv2.dilate(images[idx][1],kernel1,iterations=1), (28, 28))
        plt.imshow(image, 'gray')  # Add every number to the plot
        kernel2 = np.ones((2, 2), np.uint8)
        cv2.imwrite('./pre/class1/number%d.png' % (count), 255-cv2.dilate(image,kernel2,iterations=1))
        count += 1
    plt.savefig('./fenge/fenge.jpg', bbox_inches='tight')
    plt.close()

